import numpy as np
import math


# 球坐标转空间直角坐标
# def transformSPH2XYZ(row):
#     a = math.radians(row[0])  # 水平转角
#     b = math.radians(row[1])  # 竖直转角
#     r = row[2]  # 距离
#     x = r * math.sin(b) * math.cos(a)
#     y = -r * math.sin(b) * math.sin(a)
#     z = r * math.cos(b)
#     return [x, y, z]


def transformSPH2XYZ(row):
    a = np.deg2rad(row[:, 0])  # 水平转角
    b = np.deg2rad(row[:, 1])  # 竖直转角
    r = row[:, 2]  # 距离
    xyz = np.zeros(row.shape)
    xyz[:, 0] = r * np.sin(b) * np.cos(a)
    xyz[:, 1] = -r * np.sin(b) * np.sin(a)
    xyz[:, 2] = r * np.cos(b)
    return xyz


# 空间直角坐标转球坐标
# 适用于垂直转角大于0时
# def transformXYZ2SPH(row):
#     x = row[0]
#     y = row[1]
#     z = row[2]
#     a = np.arctan2(-y, x)
#     if a < 0:
#         a = a + math.pi * 2
#     b = 0.5 * math.pi - np.arctan2(z, (x**2 + y**2)**0.5)
#     r = z / math.cos(b)
#     return [math.degrees(a), math.degrees(b), r]


def transformXYZ2SPH(row):
    x = row[:, 0]
    y = row[:, 1]
    z = row[:, 2]
    a = np.arctan2(-1 * y, x)
    for i in range(len(a)):
        if a[i] < 0:
            a[i] = a[i] + math.pi * 2
    b = np.ones(x.shape) * 0.5 * math.pi - np.arctan2(z, (x * x + y * y)**0.5)
    r = z / np.cos(b)
    abr = np.zeros(row.shape)
    abr[:, 0] = np.rad2deg(a)
    abr[:, 1] = np.rad2deg(b)
    abr[:, 2] = r
    return abr


# 空间直角转球坐标
# 适用于垂直转角小于0时
# def transformXYZ2SPH2(row):
#     x = row[0]
#     y = row[1]
#     z = row[2]
#     a = np.arctan2(-y, x) - math.pi
#     if a < 0:
#         a = a + math.pi * 2
#     b = np.arctan2(z, (x**2 + y**2)**0.5) - 0.5 * math.pi
#     r = z / math.cos(b)
#     return [math.degrees(a), math.degrees(b), r]


def transformXYZ2SPH2(row):
    x = row[:, 0]
    y = row[:, 1]
    z = row[:, 2]
    a = np.arctan2(-1 * y, x) - np.ones(x.shape) * math.pi
    for i in range(len(a)):
        if a[i] < 0:
            a[i] = a[i] + math.pi * 2
    b = np.arctan2(z, (x * x + y * y)**0.5) - 0.5 * math.pi * np.ones(x.shape)
    r = z / np.cos(b)
    abr = np.zeros(row.shape)
    abr[:, 0] = np.rad2deg(a)
    abr[:, 1] = np.rad2deg(b)
    abr[:, 2] = r
    return abr
